import os, sys
from typing import Dict

parent_path = os.path.abspath(os.path.join(__file__, *(['..'] * 4)))
sys.path.insert(0, parent_path)

from base64 import b64encode
import traceback
import numpy as np
import socket, json, cv2, yaml, time
from typing import *
from core.task.modules.processors.base_processor import BaseProcessor
from core.utils.PIL_draw import pic_text
from core.utils.visualize import get_color_map_list
from collections import deque

class RTMPProcessor(BaseProcessor):
    """直播流检测后处理方案
    """
    save_idx = 0

    def _init_check(self):
        pass

    def _call(self, data:Dict) -> Dict:
        img_data = data[self.keys["in"]]
        det_result = data[self.keys["det"]]
        kpt_result = data[self.keys["kpt"]]
        
        for img_data, bboxes, kpts in zip(img_data, det_result["boxes"],kpt_result):
            img = img_data.copy()
            draw_thickness = min(img.shape[:2]) // 320
            for box in bboxes:
                cls_id, score, x1,y1,x2,y2 = box
                if cls_id == 0:
                    save_img = img[int(y1):int(y2),int(x1):int(x2),:]
                    cv2.imwrite("data/output/{}.jpg".format(self.save_idx),save_img)
                    self.save_idx+=1
                cls_id = int(cls_id)
                img = cv2.rectangle(img,(int(x1), int(y1)),(int(x2),int(y2)),color=[0,0,255],thickness=draw_thickness)
            for obj_kpts in kpts:
                hood = obj_kpts[3]
                if hood[2]<0.5:
                    break
                tail = obj_kpts[4]
                hood_point = hood[:2].astype(np.int32)
                img=cv2.line(img, tail[:2].astype(np.int32), hood_point,(0,0,255))
                for tip in obj_kpts[:3]:
                    if tip[2]<0.3:
                        continue
                    img=cv2.line(img, tip[:2].astype(np.int32), hood_point,(255,0,255))                
                
            data["vis"] = img
            break

            

        return data
    
class RTMPTriggerProcessor(BaseProcessor):
    """直播流检测后处理方案
    """

    continue_threshold = 5
    continue_list = deque(maxlen=5)
    last_trigger_time = time.time()

    def _init_check(self):
        pass

    def _call(self, data:Dict) -> Dict:

        det_result = data[self.keys["det"]]
        need_trigger = False
        for bboxes in det_result["boxes"]:
            have_target = False
            for box in bboxes:
                cls_id, score, x1,y1,x2,y2 = box
                cls_id = int(cls_id)
                if cls_id in [1,2,3]:
                    have_target = True
                    break
            self.continue_list.append(have_target)
            conti = self.continue_list.count(1)
            if time.time() - self.last_trigger_time>=20:
                if conti>=self.continue_threshold:
                    need_trigger = True
                    self.last_trigger_time = time.time()
        data[self.keys["trigger_tag"]]=need_trigger
            

        return data
    
    